"""
@author: deyongww
@time: 2022-8-31
algoritm 代码实现CART回归树算法模型
slogan 日拱一卒已近不行了，得日拱多卒
postscript: 持续优化
"""

"""
   所有的二叉树都有节点，所以我们先定义基础节点
"""
import pandas as pd
from matplotlib import pyplot as plt


class Node():

    def __init__(self, best_feature_id, best_split_value):
        """
        :param best_feature_id: 最优划分节点的id
        :param best_split_value:最优划分节点的值
        """
        self.feature_id = best_feature_id  # 最佳划分值样本的id
        self.split_value = best_split_value  # 最佳划分值
        self.left_child = None  # 该节点的左孩子
        self.right_child = None  # 该节点的右孩子
        self.if_leaf = False  # 判断该节点是是否是叶子节点
        self.mean_y = None  # 该节点的y值的平均值

    def set_left_chilren(self, a_node):
        self.left_child = a_node

    def set_right_chilren(self, a_node):
        self.right_child = a_node

    def get_left_chilren(self):
        return self.left_child

    def get_right_chilren(self):
        return self.right_child

    def set_as_leaf(self):
        """
        设置当前节点为叶子节点
        :return:
        """
        self.if_leaf = True


class CartTree():
    """
    有了基础节点Node类后，我们就可以基于Node创建二叉树了
    """

    def __init__(self, a_node):
        """
        Cart树也是二叉树，二叉树有更节点
        :param a_node:
        """
        self.root = a_node
        self.current_node = a_node


"""
定义CART树回归模型，模型里面有很多颗CARTTree
"""


class CARTRegressionModel():

    def __init__(self):
        # 特征的数量，也就是数据的维度，比如x1,x2,x3则feature_num = 3
        self.feature_num = None
        self.tree = None  # model里面的树模型，在这里就是CART树。
        self.tree_size = 0  # 树的棵树，比如有20棵
        self.max_depth = 10  # 整个CART回归树模型的树深度，这里默认赋值10

    def fit(self, X, Y):
        """
        传入数据进行训练，训练的目的就是CART回归树。
        :param X: 数据集的全特特征
        :param Y: 数据集的Y标签，也就是我们要训练得到的数据标签。
        :return:
        """
        self.feature_num = len(X[0])  # X[0]特征值的数量
        id_lists = [i for i in range(len(X))]  # 获取每个特征的id值
        a_node = Node(-1, -1)  # 创建一个节点
        self.tree = CartTree(a_node)  # 创建一棵CartTree,将创建的节点作为更节点
        self.build_cart_tree(X, Y, id_lists, self.tree.current_node, True, 0)

    def calculate_MSE_value(self, X, Y, id_lists, sample_id, feature_id):
        """
        在公式人工构建CAER树和公式解析阶段，我们就知道，计算平方误差MSE是选择划分节点的关键。
        :param X: 数据集的特征
        :param Y: 数据集的Y标签，也就是我们需要的结果数据
        :param id_lists: 当前节点中的所有样本的序号，比如当前节点里面有4个样本数据，id_lists=4
        :param sample_id: 当前样本的id
        :param feature_id: 当前样本的特征id
        :return: 按照每个特征数值划分后的均方误差。
        """
        best_split_value = X[sample_id, feature_id]  # 获取样本的特征值作为节点划分值

        """
        best_split_value这个划分值确定后，将数据划分为两部分
        左边部分：left <= best_split_value
        右边部分：right > best_split_value
        定义left_id_list：存储左边部分样本数据的id值
        定义right_id_list：存储右边部分样本数据的id值
        然后我们需要分别计算左边部分和右边部分的标签值之和
        定义:left_y_sum,right_y_sum
        """
        left_id_list, right_id_list = [], []
        left_y_sum, right_y_sum = 0, 0
        """
        接下来我们循环取样本数据，将取到的样本值和划分值best_split_value做比较，
        然后确定该样本是放入左边部分还是右边部分，同时将样本的Y值相加
        """
        for id in id_lists:
            if X[id, feature_id] <= best_split_value:
                left_id_list.append(id)
                left_y_sum += Y[id]
            else:
                right_id_list.append(id)
                right_y_sum += Y[id]
        """
        for循环执行完，我们已经将样本数据根据best_split_value划分为左右两部分
        并且求得了左右两部分Y值的总和。
        接下来我们需要计算左右两部分Y值总和的平均值，
        有了平均值我们就可以分别计算左右两部分值的MSE
        """
        left_y_mean = left_y_sum / len(left_id_list) if len(left_id_list) > 0 else 0
        right_Y_mean = right_y_sum / len(right_id_list) if len(right_id_list) > 0 else 0
        """
        分别定义
            左边部分均方误差：left_mse
            右边部分均方误差：right_mse
        """
        left_mse, right_mse = 0, 0
        for left_id in left_id_list:
            left_mse += (Y[left_id] - left_y_mean) ** 2
        for right_id in right_id_list:
            right_mse += (Y[right_id] - right_Y_mean) ** 2

        total_mse = left_mse + right_mse
        return total_mse

    def build_cart_tree(self, X, Y, cunrrent_id_list, parent_node, if_left_right_child, depth):
        """
        函数用于构建cart树
        :param X: 数据集的特征
        :param Y: 数据集的标签值
        :param cunrrent_id_list: 当前节点的样本序号
        :param parent_node: 当前节点的父节点
        :param if_left_right_child: True:当前节点是左子结点 False:当前节点是右子节点
        :param depth: cart树的深度
        :return:
        """
        if len(cunrrent_id_list) == 0:
            # 如果节点中没有样本，直接返回
            return
        elif len(cunrrent_id_list) == 1 or depth > self.max_depth:
            print('因为节点中只有一个样本，所以这是一个叶子节点')
            parent_node.set_as_leaf()  # 设置该节点为叶子节点。
            mean_y = 0
            for id in cunrrent_id_list:
                mean_y += Y[id]  # 因为当前就一个叶子节点
            mean_y = mean_y / len(cunrrent_id_list)
            parent_node.mean_y = mean_y
            return
        else:
            # 定义切分的最好的特征id和切分的最佳值
            best_feature_id, best_split_value = 0, 0
            # 定义均方误差
            best_mse = -1
            for sample_id in cunrrent_id_list:
                for feature_id in range(self.feature_num):
                    # 找到最佳分割特征和取值，通过调用求取均方误差函数，获取最佳的mse值
                    this_mse = self.calculate_MSE_value(X, Y, cunrrent_id_list, sample_id, feature_id)
                    if best_mse == -1 or this_mse < best_mse:
                        best_mse = this_mse
                        # 获取到最佳分割样本的id和分割值。
                        best_feature_id = feature_id
                        best_split_value = X[sample_id, best_feature_id]
                    self.tree_size += 1  # 构建了一个树，tree_size自加

                    # 变量分别存储二叉树左子树节点中的样本id，右子树节点中的样本id
                    left_id_list, right_id_list = [], []
                    for id in cunrrent_id_list:
                        if X[id, best_feature_id] <= best_split_value:
                            left_id_list.append(id)  # 样本id存储到左子树节点中
                        else:
                            right_id_list.append(id)  # 样本id存储到右子树节点中
                    print(
                        f'这个是第{self.tree_size}个节点，深度为{depth},左子结点中的样本id{len(left_id_list)}，右子节点中的样本id{len(right_id_list)}')
                    # 构建最优的划分节点对象
                    this_node = Node(best_feature_id, best_split_value)
                    # 如果该节点是左子结点，就设置为左子节点，否则为右子节点
                    if if_left_right_child:
                        parent_node.set_left_chilren(this_node)
                    else:
                        parent_node.set_right_chilren(this_node)
                    depth += 1

    def search_in_tree(self, features, current_node):
        """

        :param features:
        :param current_node:
        :return:
        """
        # 当前节点不是叶子节点
        while current_node.if_leaf != True:
            print(current_node.feature_id, features[current_node.feature_id],
                  current_node.split_value, current_node.if_leaf)
            print(current_node.left_child, current_node.right_child)
            # 如果当前特征值<= 分割值
            if features[current_node.feature_id] <= current_node.split_value:
                a_node = current_node.get_left_chilren()  # 获取左子结点
            else:
                a_node = current_node.get_right_chilren()

            if a_node == None:
                return -1
            current_node = a_node
            return current_node.mean_y

    def predict(self, X):
        """
        传入特征数据集，预测便签值y
        :param X: 特征数据集
        :return:
        """
        pred_list = []
        for sample_id in range(X.shape[0]):
            current_node = self.tree.root.left_child
            pred = self.search_in_tree(X[sample_id], current_node)
            pred_list.append(pred)

        return pred_list


# 主函数测试
if __name__ == '__main__':
    print(pd.__version__)
    data = pd.read_excel("./ENB2012_data.xlsx")
    data = data.sample(frac=1)
    data = data.values
    X, Y = data[:, 0: -2], data[:, -2]
    plt.plot(X[:, 0])
    plt.show()

    # 训练模型
    model = CARTRegressionModel()
    model.fit(X, Y)

    # 测试
    prediction = model.predict(X)
    res = [Y[i] - prediction[i] for i in range(len(prediction))]
    plt.plot(Y)
    plt.plot(prediction)
    plt.show()
